#pragma once

#include <iostream>
#include <cassert>
using namespace std;
/*
如果需要一种查询高效且有序的数据结构，且数据的个数不改变，可以考虑AVL树
AVL树不适合需要经常修改结构的场景，因为每次插入或删除后，AVL树都要进行旋转来维持平衡
但一个结构经常修改，就不太适合。此时我们需要一种性能更优的数据结构——红黑树
*/
// 设定 bf = right - left
template <class K, class V>
struct AVLTreeNode
{
    pair<K, V> _kv;
    AVLTreeNode<K, V> *_left;
    AVLTreeNode<K, V> *_right;
    AVLTreeNode<K, V> *_parent;

    int _bf; // balance factor

    AVLTreeNode(const pair<K, V> &kv)
        : _kv(kv), _left(nullptr), _right(nullptr), _parent(nullptr), _bf(0)
    {
    }
};

template <class K, class V>
class AVLTree
{
    typedef AVLTreeNode<K, V> Node;

private:
    Node *_root = nullptr;

public:
    int _rotateCount = 0;

public:
    bool Insert(const pair<K, V> &kv)
    {
        // 1.按照二叉搜索树插入结点
        if (_root == nullptr)
        {
            _root = new Node(kv);
            return true;
        }

        Node *parent = nullptr;
        Node *cur = _root;
        while (cur)
        {
            if (cur->_kv.first < kv.first)
            {
                parent = cur;
                cur = cur->_right;
            }
            else if (cur->_kv.first > kv.first)
            {
                parent = cur;
                cur = cur->_left;
            }
            else
            {
                return false;
            }
        }

        cur = new Node(kv);
        if (parent->_kv.first < kv.first)
        {
            parent->_right = cur;
        }
        else
        {
            parent->_left = cur;
        }

        cur->_parent = parent;

        // 2.控制平衡

        // 更新平衡因子
        while (parent)
        {
            if (cur == parent->_left)
            {
                parent->_bf--;
            }
            else // if (cur == parent->_right)
            {
                parent->_bf++;
            }

            if (parent->_bf == 0)
            {
                // 更新结束
                break;
            }
            else if (parent->_bf == 1 || parent->_bf == -1)
            {
                // 继续往上更新
                cur = parent;
                parent = parent->_parent;
            }
            else if (parent->_bf == 2 || parent->_bf == -2)
            {
                // 子树不平衡了，需要旋转
                if (parent->_bf == 2 && cur->_bf == 1)
                {
                    RotateL(parent);
                }
                else if (parent->_bf == -2 && cur->_bf == -1)
                {
                    RotateR(parent);
                }
                else if (parent->_bf == 2 && cur->_bf == -1)
                {
                    RotateRL(parent);
                }
                else if (parent->_bf == -2 && cur->_bf == 1)
                {
                    RotateLR(parent);
                }

                break;
            }
            else
            {
                assert(false);
            }
        }

        return true;
    }

    void RotateL(Node *dad)
    {
        ++_rotateCount;

        Node *Grandpa = dad->_parent;
        Node *sonR = dad->_right;
        Node *sonRL = sonR->_left;

        dad->_right = sonRL;
        if (sonRL)
        {
            sonRL->_parent = dad;
        }

        sonR->_left = dad;
        dad->_parent = sonR;

        if (dad == _root)
        {
            _root = sonR;
            sonR->_parent = nullptr;
        }
        else
        {
            if (Grandpa->_left == dad)
            {
                Grandpa->_left = sonR;
            }
            else
            {
                Grandpa->_right = sonR;
            }

            sonR->_parent = Grandpa;
        }

        dad->_bf = sonR->_bf = 0;
    }

    void RotateR(Node *dad)
    {
        ++_rotateCount;

        Node *Grandpa = dad->_parent;
        Node *sonL = dad->_left;
        Node *sonLR = sonL->_right;

        dad->_left = sonLR;
        if (sonLR)
            sonLR->_parent = dad;

        sonL->_right = dad;
        dad->_parent = sonL;

        if (Grandpa == nullptr)
        {
            _root = sonL;
            sonL->_parent = nullptr;
        }
        else
        {
            if (Grandpa->_left == dad)
            {
                Grandpa->_left = sonL;
            }
            else
            {
                Grandpa->_right = sonL;
            }

            sonL->_parent = Grandpa;
        }

        dad->_bf = sonL->_bf = 0;
    }

    void RotateLR(Node *dad)
    {
        Node *sonL = dad->_left;
        Node *sonLR = sonL->_right;
        int bf = sonLR->_bf;

        RotateL(sonL);
        RotateR(dad);

        if (bf == 0)
        {
            dad->_bf = 0;
            sonL->_bf = 0;
            sonLR->_bf = 0;
        }
        else if (bf == -1)
        {
            dad->_bf = 1;
            sonL->_bf = 0;
            sonLR->_bf = 0;
        }
        else if (bf == 1)
        {
            dad->_bf = 0;
            sonL->_bf = -1;
            sonLR->_bf = 0;
        }
        else
        {
            cout << "平衡因子异常:" << dad->_kv.first << "->" << dad->_bf << endl;
            assert(false);
        }
    }

    void RotateRL(Node *dad)
    {
        Node *sonR = dad->_right;
        Node *sonRL = sonR->_left;
        int bf = sonRL->_bf;

        RotateR(dad->_right);
        RotateL(dad);

        if (bf == 0)
        {
            sonR->_bf = 0;
            sonRL->_bf = 0;
            dad->_bf = 0;
        }
        else if (bf == 1)
        {
            sonR->_bf = 0;
            sonRL->_bf = 0;
            dad->_bf = -1;
        }
        else if (bf == -1)
        {
            sonR->_bf = 1;
            sonRL->_bf = 0;
            dad->_bf = 0;
        }
        else
        {
            cout << "平衡因子异常:" << dad->_kv.first << "->" << dad->_bf << endl;
            assert(false);
        }
    }

    int Height()
    {
        return Height(_root);
    }

    int Height(Node *root)
    {
        if (root == nullptr)
            return 0;

        int leftHeight = Height(root->_left);
        int rightHeight = Height(root->_right);

        return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
    }

    bool IsBalance()
    {
        return IsBalance(_root);
    }

    bool IsBalance(Node *root)
    {
        if (root == nullptr)
            return true;

        int leftHight = Height(root->_left);
        int rightHight = Height(root->_right);

        if (rightHight - leftHight != root->_bf)
        {
            cout << "平衡因子异常:" << root->_kv.first << "->" << root->_bf << endl;
            return false;
        }

        return abs(rightHight - leftHight) < 2 && IsBalance(root->_left) && IsBalance(root->_right);
    }
    // 在删除节点后，需要更新从该节点到根节点路径上所有节点的平衡因子，
    // 并进行平衡调整，使得整棵树重新满足AVL树的性质。
    // 删除接口
    // bool Remove(const pair<K, V> &kv)；

};
